import math

from agent import Agent, Drone


class Controller:
    def __init__(self,
                 agent: 'Agent'):
        self.agent = agent
        self.waypoints = []

    def follow_trajectory(self):
        if not self.agent.freeze:
            kp_speed = 0.1  # 速度比例增益
            if self.waypoints:
                while True:
                    # 更新路径点跟踪情况
                    if not self.waypoints:
                        # 无路径点则停止
                        break
                    desired_x, desired_y, _ = self.waypoints[0]['position']
                    arrival_time = self.waypoints[0]['arrival_time']
                    error_x = desired_x - self.agent.x
                    error_y = desired_y - self.agent.y
                    distance_error = math.sqrt(error_x ** 2 + error_y ** 2)
                    if distance_error > self.agent.turn_radius:
                        break
                    else:
                        # 移除已经抵达的路径点
                        self.waypoints.pop(0)
                if self.waypoints:
                    # 有航路点则持续追踪下一个航路点
                    # 期望航向
                    self.agent.target_heading = math.atan2(error_y, error_x)
                    # 期望速度 - 简单比例
                    self.agent.target_speed = min(self.agent.max_speed,
                                                  max(
                                                      distance_error / max(1, arrival_time - self.agent.world.time),
                                                      self.agent.min_speed))
                else:
                    # 无航路点则根据智能体类型做出行为
                    if isinstance(self.agent, Drone):
                        # 无人机无航路点则原地盘旋
                        self.agent.target_speed = self.agent.max_speed
                        self.agent.target_heading = self.agent.yaw + math.pi / 2
                        pass
                    else:
                        # 无人船无航路点则停止
                        self.agent.target_speed = 0
            else:
                # 无航路点则根据智能体类型做出行为
                if isinstance(self.agent, Drone):
                    # 无人机无航路点则原地盘旋
                    self.agent.target_speed = self.agent.max_speed
                    self.agent.target_heading = self.agent.yaw + math.pi / 2
                    pass
                else:
                    # 无人船无航路点则停止
                    self.agent.target_speed = 0
            # 计算跟踪误差
            speed_error = self.agent.target_speed - self.agent.v
            heading_error = self.agent.target_heading - self.agent.yaw
            # 角度归一化
            while heading_error > math.pi:
                heading_error -= 2 * math.pi
            while heading_error < -math.pi:
                heading_error += 2 * math.pi
            # 控制输出 - 直接控制
            self.agent.cmd_accel = kp_speed * speed_error
            self.agent.cmd_accel = max(-self.agent.max_accel, min(self.agent.cmd_accel, self.agent.max_accel))
            self.agent.omega = heading_error / self.agent.world.interval
            self.agent.omega = max(-self.agent.v / (self.agent.turn_radius * (2 if self.agent.world.time < 3600 else 1)),
                                   min(self.agent.omega, self.agent.v / (self.agent.turn_radius * (2 if self.agent.world.time < 3600 else 1))))
        else:
            if isinstance(self.agent, Drone):
                pass
            else:
                # 冻结状态失去动力，减速为0
                self.agent.cmd_accel = -self.agent.max_accel
                # 冻结状态失去角速度控制能力，角速度减小为0
                self.agent.omega = 0.5 * self.agent.omega

    def fast_move_to_position(self, position):
        self.waypoints = [{'position': position, 'arrival_time': 0}]

    def set_speed(self, velocity:float, direction:float):
        if self.agent.world is not None:
            time = self.agent.world.time + 1000000 / velocity
        else:
            time = 1000000 / velocity
        self.waypoints = [{'position': (self.agent.x + 1000000 * math.cos(direction), self.agent.y + 1000000 * math.sin(direction), 0), 'arrival_time': time}]